package cn.metona.mq.core;

import java.io.*;
import java.nio.file.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;
import java.util.logging.Logger;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

/**
 * 消息存储管理器
 * 负责消息的内存存储和持久化
 */
public class MessageStore {
    private static final Logger logger = Logger.getLogger(MessageStore.class.getName());

    // 内存中的消息存储
    private final ConcurrentHashMap<String, QueueMessage> memoryStore;

    // 持久化文件路径
    private final Path persistencePath;

    // 消息序列号生成器
    private final AtomicLong sequenceGenerator;

    // 是否启用持久化
    private final boolean persistenceEnabled;

    // 批量持久化大小
    private final int batchPersistSize;

    // 持久化线程池
    private final ScheduledExecutorService persistenceExecutor;

    // 存储锁
    private final Object storeLock = new Object();

    /**
     * 构造函数
     */
    public MessageStore(String storeName, boolean enablePersistence) {
        this.memoryStore = new ConcurrentHashMap<>();
        this.persistencePath = Paths.get("data", storeName + ".dat");
        this.sequenceGenerator = new AtomicLong(0);
        this.persistenceEnabled = enablePersistence;
        this.batchPersistSize = 100;
        this.persistenceExecutor = Executors.newScheduledThreadPool(1, r -> {
            Thread t = new Thread(r, "MessageStore-Persistence-" + storeName);
            t.setDaemon(true);
            return t;
        });

        // 创建数据目录
        try {
            Files.createDirectories(persistencePath.getParent());
        } catch (IOException e) {
            logger.warning("Failed to create data directory: " + e.getMessage());
        }

        // 启动持久化任务
        if (persistenceEnabled) {
            startPersistenceTask();
        }
    }

    /**
     * 存储消息
     */
    public void storeMessage(QueueMessage message) {
        if (message == null) {
            throw new IllegalArgumentException("Message cannot be null");
        }
        memoryStore.put(message.getMessageId(), message);
        sequenceGenerator.incrementAndGet();
    }

    /**
     * 获取消息
     */
    public QueueMessage getMessage(String messageId) {
        if (messageId == null) {
            return null;
        }
        return memoryStore.get(messageId);
    }

    /**
     * 删除消息
     */
    public void removeMessage(String messageId) {
        if (messageId != null) {
            memoryStore.remove(messageId);
        }
    }

    /**
     * 获取所有消息ID
     */
    public Set<String> getAllMessageIds() {
        return new HashSet<>(memoryStore.keySet());
    }

    /**
     * 获取消息数量
     */
    public long getMessageCount() {
        return memoryStore.size();
    }

    /**
     * 清空所有消息
     */
    public void clear() {
        memoryStore.clear();
        sequenceGenerator.set(0);
    }

    /**
     * 启动持久化任务
     */
    private void startPersistenceTask() {
        persistenceExecutor.scheduleAtFixedRate(() -> {
            try {
                persistMessages();
            } catch (Exception e) {
                logger.severe("Error during message persistence: " + e.getMessage());
            }
        }, 30, 30, TimeUnit.SECONDS); // 每30秒持久化一次
    }

    /**
     * 持久化消息到文件
     */
    public void persistMessages() {
        if (!persistenceEnabled || memoryStore.isEmpty()) {
            return;
        }

        synchronized (storeLock) {
            try {
                // 创建临时文件
                Path tempPath = persistencePath.resolveSibling(persistencePath.getFileName() + ".tmp");

                // 使用try-with-resources确保资源正确关闭
                try (ObjectOutputStream oos = new ObjectOutputStream(
                        new GZIPOutputStream(Files.newOutputStream(tempPath)))) {

                    // 序列化消息数据
                    oos.writeObject(new HashMap<>(memoryStore));
                    oos.writeObject(sequenceGenerator.get());

                    // 原子性替换文件
                    Files.move(tempPath, persistencePath,
                            StandardCopyOption.REPLACE_EXISTING,
                            StandardCopyOption.ATOMIC_MOVE);

                    logger.fine("Messages persisted to disk: " + memoryStore.size() + " messages");
                } catch (IOException e) {
                    // 如果临时文件创建失败，尝试删除
                    try {
                        Files.deleteIfExists(tempPath);
                    } catch (IOException deleteEx) {
                        logger.warning("Failed to delete temp file: " + deleteEx.getMessage());
                    }
                    throw e;
                }
            } catch (Exception e) {
                logger.severe("Failed to persist messages: " + e.getMessage());
            }
        }
    }

    /**
     * 从文件恢复消息
     */
    @SuppressWarnings("unchecked")
    public void recoverMessages() {
        if (!persistenceEnabled || !Files.exists(persistencePath)) {
            return;
        }

        synchronized (storeLock) {
            try (ObjectInputStream ois = new ObjectInputStream(
                    new GZIPInputStream(Files.newInputStream(persistencePath)))) {

                // 反序列化消息数据
                Object messagesObj = ois.readObject();
                Object sequenceObj = ois.readObject();

                if (messagesObj instanceof HashMap && sequenceObj instanceof Long recoveredSequence) {
                    HashMap<String, QueueMessage> recoveredMessages =
                            (HashMap<String, QueueMessage>) messagesObj;

                    // 恢复到内存
                    memoryStore.clear();
                    memoryStore.putAll(recoveredMessages);
                    sequenceGenerator.set(recoveredSequence);

                    logger.info("Messages recovered from disk: " + recoveredMessages.size() + " messages");
                }
            } catch (Exception e) {
                logger.severe("Failed to recover messages: " + e.getMessage());
                // 如果恢复失败，删除损坏的文件
                try {
                    Files.deleteIfExists(persistencePath);
                } catch (IOException deleteEx) {
                    logger.warning("Failed to delete corrupted persistence file: " + deleteEx.getMessage());
                }
            }
        }
    }

    /**
     * 关闭存储
     */
    public void close() {
        if (persistenceEnabled) {
            // 最后一次持久化
            persistMessages();
        }

        persistenceExecutor.shutdown();
        try {
            if (!persistenceExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
                persistenceExecutor.shutdownNow();
            }
        } catch (InterruptedException e) {
            persistenceExecutor.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }

    /**
     * 获取下一个序列号
     */
    public long getNextSequence() {
        return sequenceGenerator.incrementAndGet();
    }

    public int getBatchPersistSize() {
        return batchPersistSize;
    }
}